import cnn
import numpy
import datastore
import sys
import argparse
import torch
import dataaugment
import random
import time
import os
import datastore
#from torch.utils.tensorboard import SummaryWriter
#tb_writer = SummaryWriter()

parser = argparse.ArgumentParser(description='vslidate model.')
parser.add_argument("--src", required=True)

args = parser.parse_args()

in_path = args.src

classes = datastore.load_classes(in_path)

TTrain = datastore.audioDatastore(classes)
TTrain.add_dataLoc(os.path.join(args.src,"validation"))

numHops = 98

numClasses = len(classes)
testloader = TTrain.gen_dataset()

model = cnn.ConvNeuralNet(numHops, numClasses)
model.load_state_dict(torch.load("speech-cmd-model.pth"))
model.eval()

acc = 0
count = 0
stat=[0]*numClasses

for n in range(len(testloader)):
    i, j = testloader[n]
    # spectrum = TTrain.load(i, j)
    data = TTrain.load(i, j)

    # data,_ = dataaugment.add_s_p_noise(spectrum, c, points[n])
    inputs = torch.from_numpy(numpy.asarray([[data]], dtype=numpy.float32))
    labels = torch.from_numpy(numpy.asarray([i]))
    y_pred = model(inputs)

    #print(y_pred)
    #print(torch.argmax(y_pred, 1) )
    # print(torch.argmax(y_pred, 1) == labels)
    argmax=torch.argmax(y_pred, 1)
    #print(argmax.item())
    if argmax.item() == i:
        stat[i] +=1
        acc +=1
    #acc += ( argmax== labels ).float().sum()
    count += 1  # len(labels)

percent = acc / count
print("model accuracy %.2f%% %d %d" % ( percent * 100, count, acc))
print(stat)
